"""
Created on Sat Apr 22 14:27:23 2023

@author: shosborn
"""

#!/usr/bin/python3

import SMART
import numpy as np
from random import *
from multiprocessing.pool import ThreadPool
import sys

OBLIVIOUS = SMART.Partition.OBLIVIOUS
AWARE_START_PHYS = SMART.Partition.AWARE_START_PHYS
AWARE_START_THREAD = SMART.Partition.AWARE_START_THREAD
AWARE_START_OBLIV = SMART.Partition.AWARE_START_OBLIV
OPTIMAL = SMART.Partition.OPTIMAL
NUM_METHODS = 5

core_count = int(sys.argv[1])
FILE_OUT = sys.argv[2]
# FILE_OUT = "srt_runTests_test.csv"

# DISSERTATION might be a better way to do this

# returns theBins
# What is theBins?
# [taskSystemsAttempted][schedulable on m cores via method 1, sched via method 2, etc.]
# first item in outer list gives number of task systems attempted.
# subsequent items correspond to different utilization ranges
# each subsequent item is a list with each item in the inner list giving the number of schedulable systmes for the
# method corresponding to that items location
def run_param_test(core_count, bin_width, sample_target, begin_util, per_task_util_min, per_task_util_max,
                   expected_score, prob_harm, degree_harm, score_distro, useOpt=0):
    # create enough bins to go up to twice the core count
    numBins = 2 * int(np.ceil(core_count / bin_width))
    # "theBins": index 0 is number of task sets in bin, following indicies are how many of those are schedulable for various algorithms
    theBins = []
    for i in range(0, numBins):
        # has 11 because results include extra testing methods that got dropped.
        theBins.append([0] * 11)

    parameterResults = []

    # Iterate until we've tried at least 'sample_target' task sets at our base utilization
    # Each iteration of this tests at least one task set per bin
    while theBins[0][0] <= sample_target:
        myTaskSet = SMART.TaskSet(begin_util, per_task_util_min, per_task_util_max, expected_score, prob_harm,
                                  degree_harm, score_distro)
        # Loop until we reach a task set that nothing can schedule
        while True:
            # Make sure we have at least enough utilization to reach the first bin (at m)
            # we're not testing util < core_count; those will always succeed
            while myTaskSet.totalUtil < core_count:
                myTaskSet.addTask()
            # Depending on where our utilization fell as we added tasks, we choose a bin
            myBin = int(np.floor((myTaskSet.totalUtil - core_count) / bin_width))
            theBins[myBin][0] += 1  # One more sample for this bin
            # Test if each algorthm can schedule this set
            # "optimal" gets used only if useOpt = 1 (it defaults to 0)
            for method in range(OBLIVIOUS, OPTIMAL + useOpt):
                myTaskSet.partitionTasks(method)
                theseResults = myTaskSet.partitionList[method].testUmaFunk(core_count)
                theBins[myBin][method] += theseResults[SMART.UMA_RESULT]
                theBins[myBin][method + NUM_METHODS] += theseResults[SMART.FUNK_RESULT]
            # If any algorithm had any successes, keep adding work (and potentially fall into higher bins)
            # until that's not true
            if sum(theBins[myBin][OBLIVIOUS:]):
                myTaskSet.addTask()
            else:
                break
                # If we reach a task set that's unschedulable, assume it's unschedulable for higher utilizations.
        # This step records that assumption for bins after the one that everyone failed on
        # (We could just simulate all the higher-utilization bins, but that's slow.)
        totalUtil = myTaskSet.totalUtil
        while totalUtil < core_count:
            totalUtil += random() * (myTaskSet.utilMax - myTaskSet.utilMin) + myTaskSet.utilMin
            myBin = int(np.floor((totalUtil - core_count) / bin_width))
            theBins[myBin][0] += 1
    return theBins


def calculate_rsa(theBins, algo, bin_width, core_count):
    # total_samples = theBins[0]
    total_area = core_count
    # iterate through theBins
    for b in theBins:
        if b[0] > 0:
            total_area += (b[algo] / b[0]) * bin_width
    return total_area / core_count


# compared to run_param test, also needs theBins
# begin_util, and sample_target aren't used, but putting them in signature allows run_test to just use
# config as argument
# does not need useOpt
def save_scenario_results(theBins, core_count, begin_util, per_task_util_min, per_task_util_max,
                          expected_score, prob_harm,
                          degree_harm, score_distro, bin_width, sample_target):
    # build the parameter string
    parameter_str = ",".join(
        [str(core_count), str(per_task_util_min), str(per_task_util_max), str(expected_score), str(prob_harm),
         str(degree_harm), str(score_distro)])
    # for each algorithm:
    ALGO_INDEX = 0
    RSA_INDEX = 1
    max_rsa = 0
    max_rsa_algo = -1
    # setting up this way so not reliant on the algo indices being consecutive starting from zero
    rsa_list = [0] * (max(OBLIVIOUS, AWARE_START_PHYS, AWARE_START_THREAD, AWARE_START_OBLIV) + 1)
    for algo in (OBLIVIOUS, AWARE_START_PHYS, AWARE_START_THREAD, AWARE_START_OBLIV):
        rsa_list[algo] = calculate_rsa(theBins, algo, bin_width, core_count)
        # is it the best?
        if rsa_list[algo] > max_rsa:
            max_rsa = rsa_list[algo]
            max_rsa_algo = algo
    for algo in (OBLIVIOUS, AWARE_START_PHYS, AWARE_START_THREAD, AWARE_START_OBLIV):
        output_string = ",".join([parameter_str, str(algo), str(rsa_list[algo])])
        # was this algorithm the best?
        # would it be better to list the best algo and rsa?
        # can change how exactly I do this to make the graph easier.
        if max_rsa_algo == algo:
            output_string = output_string + "," + str("1")
        else:
            output_string = output_string + "," + str("0")
    # add the actual results to the string
    # repeated concatenation is really inefficient
        for b in theBins:
            if b[0] > 0:
                output_string = output_string + "," + str(b[algo] / b[0])
            else:   # prevent divide-by-zero
                output_string = output_string + ",0"
        output_string = output_string + ","

        with open(FILE_OUT, "a") as f:
            print(output_string, file=f)
        # print("Complete: ", parameter_str)


# come up with some test parameters
# call run_param_test
# print results to file

'''
Whatever is covered by this (nested) loop will come out as one file.

'''

# real parameters
# try different bin size values before starting big run
# not sure if this should be absolute or relative to core count

BIN_SIZE = 0.10 # If too small, will have uneven samples/bin distribution. If too large, will take forever + bad RSA
SAMPLE_TARGET = 1000 # try different values before starting the big run
UTIL_RANGES = ((0, 0.4), (0.3, 0.7), (0.6, 1), (0, 1))
EXPECTED_SCORES = (0.2, 0.4, 0.6)
PROB_HARM = (0, 0.125, 0.25)
SCORE_DISTRO = (SMART.FIXED, SMART.EXPO)

'''
# Test Parameters
BIN_SIZE = 0.10
SAMPLE_TARGET = 10  # try different values before starting the big run
# CORE_COUNTS = (4, 8)
UTIL_RANGES = ((0, 0.4), (0.6, 1))
EXPECTED_SCORES = [0.2, 0.4]
PROB_HARM = [0, 0.125]
SCORE_DISTRO = (SMART.FIXED, SMART.EXPO)
core_count = 4
'''
# core_count is an input parameter
configurations = []
for util_range in UTIL_RANGES:
    for expected_score in EXPECTED_SCORES:
        for prob_harm in PROB_HARM:
            for score_distro in SCORE_DISTRO:
                config = {}
                config["core_count"] = core_count
                config["begin_util"] = core_count
                config["bin_width"] = BIN_SIZE
                config["sample_target"] = SAMPLE_TARGET
                config["per_task_util_min"] = util_range[0]
                config["per_task_util_max"] = util_range[1]
                config["expected_score"] = expected_score
                config["prob_harm"] = prob_harm
                config["degree_harm"] = 2
                config["score_distro"] = score_distro
                configurations.append(config)

# Run mode configuration
SYNC = False


# the input to run_test, 'config', is one of the many dictionaries contained in the list 'configurations'
def run_test(config):
    # the '**' notation indicates that dictionary config is to be treated as a set of keyword arguments
    theBins = run_param_test(**config)
    # run_param_test and print_scenario_results need to have consistent parameter names
    save_scenario_results(theBins, **config)

# We need this protected so that subprocesses do not run it
if __name__ == '__main__':
    with open(FILE_OUT, "a") as f:
        #print headers
        print(
            "core_count,per_task_util_min,per_task_util_max,expected_score,prob_harm,degree_harm,score_distro,algo,rsa,best_rsa,results,",
            file=f)

    # Use multithreading

    with ThreadPool(processes=len(configurations)) as pool:
        # If we do not need the results to be in order, save things whenever we get them
        for res_tup in pool.imap_unordered(run_test, configurations):
            pass
            # each result gives data corresponding to one graph

